package com.bujunjie.study.realtime.common.function;

import com.alibaba.fastjson.JSONObject;
import com.bujunjie.study.realtime.common.constant.FlinkConstant;
import com.bujunjie.study.realtime.common.util.HBaseUtil;
import com.bujunjie.study.realtime.common.util.RedisUtil;
import io.lettuce.core.api.StatefulRedisConnection;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;
import org.apache.hadoop.hbase.client.AsyncConnection;

import java.util.Collections;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

/**
 * <p>dim 异步查询数据</p>
 *
 * @author bu.junjie
 * @version 1.0.0
 * @createTime 2025/9/18 17:07
 */
@Slf4j
public abstract class DimAsyncFunction<T> extends RichAsyncFunction<T, T> implements DimJoinFunction<T> {

    private AsyncConnection hbaseAsyncConn;
    private StatefulRedisConnection<String, String> redisAsyncConn;

    @Override
    public void open(Configuration parameters) throws Exception {
        hbaseAsyncConn = HBaseUtil.getHBaseAsyncConnection();
        redisAsyncConn = RedisUtil.getRedisAsyncConnection();
    }

    @Override
    public void close() throws Exception {
        HBaseUtil.closeAsyncHbaseConnection(hbaseAsyncConn);
        RedisUtil.closeRedisAsyncConnection(redisAsyncConn);
    }

    @Override
    public void asyncInvoke(T obj, ResultFuture<T> resultFuture) throws Exception {
        //  异步执行数据
        CompletableFuture.supplyAsync(new Supplier<JSONObject>() {
                    @Override
                    public JSONObject get() {
                        String rowKey = getRowKey(obj);
                        //根据维度的主键到Redis中获取维度数据
                        return RedisUtil.readDimAsync(redisAsyncConn, getTableName(), rowKey);
                    }
                })
                .thenApply(new Function<JSONObject, JSONObject>() {
                    @Override
                    public JSONObject apply(JSONObject dimJsonObj) {
                        if (dimJsonObj != null) {
                            //如果在Redis中找到了要关联的维度(缓存命中)。 直接将命中的维度作为结果返回
                            log.info("~~~从Redis中找到了TableName=【{}】表的 rowKey=【{}】数据~~~", getTableName(), getRowKey(obj));
                        } else {
                            //如果在Redis中没有找到要关联的维度，发送请求到HBase中查找
                            dimJsonObj = HBaseUtil.readDimAsync(hbaseAsyncConn, FlinkConstant.HBASE_NAMESPACE, getTableName(), getRowKey(obj));
                            //将查找到的维度数据放到Redis中缓存起来，方便下次查询使用
                            if (dimJsonObj != null) {
                                log.info("~~~从HBase中找到了TableName=【{}】表的 rowKey=【{}】数据~~~", getTableName(), getRowKey(obj));
                                RedisUtil.writeDimAsync(redisAsyncConn, getTableName(), getRowKey(obj), dimJsonObj);
                            } else {
                                log.info("~~~没有找到 TableName=【{}】表的 rowKey=【{}】数据~~~", getTableName(), getRowKey(obj));
                            }
                        }
                        return dimJsonObj;
                    }
                })
                .thenAcceptAsync(
                        new Consumer<JSONObject>() {
                            @Override
                            public void accept(JSONObject dimJsonObj) {
                                if (dimJsonObj != null) {
                                    //将维度对象相关的维度属性补充到流中对象上
                                    addDims(obj, dimJsonObj);
                                }
                                //获取数据库交互的结果并发送给ResultFuture的回调函数，将关联后的数据传递到下游
                                resultFuture.complete(Collections.singleton(obj));
                            }
                        }
                );

    }


}
